from argparse import Namespace
from datam import DataM
from    torch.utils.data import  DataLoader
from pokemon import Pokemon

config = Namespace(
project_name="pokeman-自定义ds-ImageFolder",
    # AVAIL_GPUS = min(1, torch.cuda.device_count()),
AVAIL_GPUS = 1,
    # BATCH_SIZE = 256 if min(1, torch.cuda.device_count()) else 64,
BATCH_SIZE = 64,
    lr = 0.02,
    optim_type = "Adam",
    max_epochs = 3,
    ckpt_path = "best.ckpt",
    data_dir ="../../../../data/pokeman",
    dropout_p = 0.1,
    modelPATH = "modelPATH",
num_workers = 2,
resize = 224,


)

def getData():
    data_mnist = DataM(data_dir=config.data_dir, batch_size=config.BATCH_SIZE, num_workers=config.num_workers,
                       resize=config.resize)
    data_mnist.setup()
    train_loader = data_mnist.train_dataloader()
    val_loader = data_mnist.val_dataloader()
    test_loader = data_mnist.test_dataloader()

    return train_loader,val_loader,test_loader


def getPolemon():
    train_db = Pokemon(config.data_dir, config.resize, mode='train')
    val_db = Pokemon(config.data_dir, config.resize, mode='val')
    test_db = Pokemon(config.data_dir, config.resize, mode='test')
    train_loader = DataLoader(train_db, batch_size=config.BATCH_SIZE, shuffle=True,
                              num_workers=config.num_workers)
    val_loader = DataLoader(val_db, batch_size=config.BATCH_SIZE, num_workers=config.num_workers)
    test_loader = DataLoader(test_db, batch_size=config.BATCH_SIZE, num_workers=config.num_workers)
    return train_loader,val_loader,test_loader




